package chapter11

import org.apache.spark.sql.expressions.Aggregator
import org.apache.spark.sql.Encoder
import org.apache.spark.sql.Encoders
import org.apache.spark.sql.SparkSession

/**
 * author: yuhui
 * descriptions: 使用强类型用户自定义UDAF入门示例：求薪资的平均值
 * date: 2024 - 11 - 29 2:14 下午
 */

// 既然是强类型，可能有case类
case class Employee(name: String, salary: Long)

case class Average(var sum: Long, var count: Long)

object UDAFStrong extends Aggregator[Employee, Average, Double] {
  // 定义一个数据结构，保存工资总数和工资总个数，初始都为0
  def zero: Average = Average(0L, 0L)

  // Combine two values to produce a new value. For performance, the function may modify `buffer`
  // and return it instead of constructing a new object
  def reduce(buffer: Average, employee: Employee): Average = {
    buffer.sum += employee.salary
    buffer.count += 1
    buffer
  }

  // 聚合不同execute的结果
  def merge(b1: Average, b2: Average): Average = {
    b1.sum += b2.sum
    b1.count += b2.count
    b1
  }

  // 计算输出
  def finish(reduction: Average): Double = reduction.sum.toDouble / reduction.count

  // 设定之间值类型的编码器，要转换成case类
  // Encoders.product是进行scala元组和case类转换的编码器
  def bufferEncoder: Encoder[Average] = Encoders.product

  // 设定最终输出值的编码器
  def outputEncoder: Encoder[Double] = Encoders.scalaDouble
}


object UDAFStrongMain {

  def main(args: Array[String]): Unit = {

    val spark: SparkSession = SparkSession
      .builder()
      .appName("")
      .master("local[*]")
      .getOrCreate()

    import spark.implicits._

    val ds =  spark.sparkContext.parallelize(
      List(
        Employee("余辉",20000),
        Employee("视频号：辉哥大数据",30000),
        Employee("抖音：辉哥大数据",40000))
    ).toDS()

    ds.show()
//    +------------------+------+
//    |              name|salary|
//    +------------------+------+
//    |              余辉| 20000|
//    |视频号：辉哥大数据  | 30000|
//    |  抖音：辉哥大数据  | 40000|
//    +------------------+------+

    // Convert the function to a `TypedColumn` and give it a name
    val averageSalary = UDAFStrong.toColumn.name("avgsalary")
    val result = ds.select(averageSalary)
    result.show()
//    +--------------+
//    |avgsalary     |
//    +--------------+
//    |       30000.0|
//    +--------------+

  }
}

